package com.btl.component.crypto.util;

import com.btl.component.crypto.dto.CryptoPrivateKey;
import com.btl.component.crypto.dto.CryptoPublicKey;
import lombok.extern.slf4j.Slf4j;
import org.springframework.core.io.Resource;
import org.springframework.core.io.support.PathMatchingResourcePatternResolver;
import org.springframework.core.io.support.ResourcePatternResolver;
import org.springframework.util.StopWatch;

import javax.crypto.Cipher;
import java.io.IOException;
import java.io.InputStream;
import java.nio.charset.StandardCharsets;
import java.security.*;
import java.security.cert.CertificateException;
import java.security.cert.X509Certificate;
import java.time.ZoneId;
import java.util.Base64;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import java.util.concurrent.ConcurrentHashMap;

@Slf4j
public class RsaCryptoUtil {

    private static final Map<String, CryptoPublicKey> publicKeyMap = new ConcurrentHashMap<>();

    private static final Map<String, CryptoPrivateKey> privateKeyMap = new ConcurrentHashMap<>();

    private RsaCryptoUtil() {
    }

    /**
     * 获取公钥
     *
     * @param keyName 加密算法
     */
    public static PublicKey getPublicKey(String keyName, String keyStorePath, String keyStorePassword) throws IOException {
        CryptoPublicKey cryptoPublicKey = publicKeyMap.get(keyName + keyStorePath);
        if (Objects.isNull(cryptoPublicKey)) {
            ResourcePatternResolver resolver = new PathMatchingResourcePatternResolver();
            Resource[] resources = resolver.getResources(keyStorePath);
            for (Resource resource : resources) {
                String filename = resource.getFilename();
                String fileKeyName = Optional.ofNullable(filename)
                        .orElse("")
                        .substring(0, Optional.ofNullable(filename)
                                .orElse("")
                                .lastIndexOf("."));
                if (keyName.equalsIgnoreCase(fileKeyName)) {
                    CryptoPublicKey cryptoPublicKeyInit = getPk12ForPublicKey(resource, keyStorePassword);
                    publicKeyMap.put(keyName + keyStorePath, cryptoPublicKeyInit);
                    log.debug("当前使用公钥端口: {},  publicKey: {}", keyName + keyStorePath, Base64.getEncoder()
                            .encodeToString(cryptoPublicKeyInit.getPublicKey()
                                    .getEncoded()));
                    return cryptoPublicKeyInit.getPublicKey();
                }
            }
        }
        log.debug("当前使用公钥端口: {},  publicKey: {}", keyName + keyStorePath, Base64.getEncoder()
                .encodeToString(cryptoPublicKey.getPublicKey()
                        .getEncoded()));
        return cryptoPublicKey.getPublicKey();
    }

    /**
     * 获取私钥
     *
     * @param keyName 加密算法
     */
    public static PrivateKey getPrivateKey(String keyName, String keyStorePath, String keyStorePassword) throws IOException {
        CryptoPrivateKey cryptoPrivateKey = privateKeyMap.get(keyName + keyStorePath);
        if (Objects.isNull(cryptoPrivateKey)) {
            ResourcePatternResolver resolver = new PathMatchingResourcePatternResolver();
            Resource[] resources = resolver.getResources(keyStorePath);
            for (Resource resource : resources) {
                String filename = resource.getFilename();
                String fileKeyName = Optional.ofNullable(filename)
                        .orElse("")
                        .substring(0, Optional.ofNullable(filename)
                                .orElse("")
                                .lastIndexOf("."));
                if (keyName.equalsIgnoreCase(fileKeyName)) {
                    CryptoPrivateKey pk12ForPrivateKey = getPk12ForPrivateKey(resource, keyStorePassword);
                    privateKeyMap.put(keyName + keyStorePath, pk12ForPrivateKey);
                    log.debug("当前使用私钥端口: {},  privateKey: {}", keyName + keyStorePath, Base64.getEncoder()
                            .encodeToString(pk12ForPrivateKey.getPrivateKey()
                                    .getEncoded()));
                    return pk12ForPrivateKey.getPrivateKey();
                }
            }
        }
        log.debug("当前使用私钥端口: {},  privateKey: {}", keyName + keyStorePath, Base64.getEncoder()
                .encodeToString(cryptoPrivateKey.getPrivateKey()
                        .getEncoded()));
        return cryptoPrivateKey.getPrivateKey();
    }

    private static CryptoPrivateKey getPk12ForPrivateKey(Resource resource, String keyStorePassword) {
        try (InputStream is = resource.getInputStream()) {
            KeyStore keyStore = KeyStore.getInstance("PKCS12");
            keyStore.load(is, keyStorePassword.toCharArray());
            String alias = keyStore.aliases()
                    .nextElement();
            PrivateKey key = (PrivateKey) keyStore.getKey(alias, keyStorePassword.toCharArray());
            String filename = resource.getFilename();
            CryptoPrivateKey cryptoPrivateKey = new CryptoPrivateKey();
            cryptoPrivateKey.setFilename(filename);
            cryptoPrivateKey.setPrivateKey(key);
            return cryptoPrivateKey;
        } catch (UnrecoverableKeyException | IOException | NoSuchAlgorithmException | CertificateException |
                 KeyStoreException e) {
            throw new RuntimeException(e);
        }
    }

    private static CryptoPublicKey getPk12ForPublicKey(Resource res, String keyStorePassword) {
        try (InputStream is = res.getInputStream()) {
            KeyStore ks = KeyStore.getInstance("PKCS12");
            ks.load(is, keyStorePassword.toCharArray());
            String alias = ks.aliases()
                    .nextElement();
            String filename = res.getFilename();
            X509Certificate cert = (X509Certificate) ks.getCertificate(alias);
            CryptoPublicKey cryptoPublicKey = new CryptoPublicKey();
            cryptoPublicKey.setName(filename);
            cryptoPublicKey.setPublicKey(cert.getPublicKey());
            cryptoPublicKey.setExpireTime(cert.getNotAfter()
                    .toInstant()
                    .atZone(ZoneId.systemDefault())
                    .toLocalDateTime());
            return cryptoPublicKey;
        } catch (CertificateException | NoSuchAlgorithmException | IOException | KeyStoreException e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * 加密
     *
     * @param keyName 密钥名称
     * @param data    待加密数据
     */
    public static String encrypt(String keyName, String keyStorePath, String keyStorePassword, String data) throws Exception {
        StopWatch stopWatch = new StopWatch("rsa-encrypt");
        stopWatch.start();
        Cipher cipher = Cipher.getInstance("RSA/ECB/PKCS1Padding");
        cipher.init(Cipher.ENCRYPT_MODE, getPublicKey(keyName, keyStorePath, keyStorePassword));
        String encoded = Base64.getEncoder()
                .encodeToString(cipher.doFinal(data.getBytes(StandardCharsets.UTF_8)));
        stopWatch.stop();
        log.debug("rsa-encrypt: {}ms", stopWatch.getTotalTimeMillis());
        return encoded;
    }

    /**
     * 解密
     *
     * @param encryptedData 待解密数据
     */
    public static String decrypt(String vName, String keyStorePath, String keyStorePassword, String encryptedData) throws Exception {
        StopWatch stopWatch = new StopWatch("rsa-decrypt");
        stopWatch.start();
        byte[] decode = Base64.getDecoder()
                .decode(encryptedData);
        Cipher cipher = Cipher.getInstance("RSA/ECB/PKCS1Padding");
        cipher.init(Cipher.DECRYPT_MODE, getPrivateKey(vName, keyStorePath, keyStorePassword));
        byte[] decryptedBytes = cipher.doFinal(decode);
        String decrypted = new String(decryptedBytes, StandardCharsets.UTF_8);
        stopWatch.stop();
        log.debug("rsa-decrypt: {}ms", stopWatch.getTotalTimeMillis());
        return decrypted;
    }

    /**
     * 生成签名
     *
     * @param data  待签名数据
     * @param vName 密钥名称
     * @return Base64编码的签名结果
     */
    public static String sign(String vName, String keyStorePath, String keyStorePassword, String data) throws Exception {
        StopWatch stopWatch = new StopWatch("rsa-sign");
        stopWatch.start();
        Signature signature = Signature.getInstance("SHA256WithRSA");
        signature.initSign(getPrivateKey(vName, keyStorePath, keyStorePassword));
        signature.update(data.getBytes(StandardCharsets.UTF_8));
        String signed = Base64.getEncoder()
                .encodeToString(signature.sign());
        stopWatch.stop();
        log.debug("rsa-sign: {}ms", stopWatch.getTotalTimeMillis());
        return signed;
    }

    /**
     * 验证签名
     *
     * @param vName 密钥名称
     * @param data  原始数据
     * @param sign  待验证的签名(Base64编码)
     * @return 验证结果
     */
    public static boolean verify(String vName, String keyStorePath, String keyStorePassword, String data, String sign) throws Exception {
        StopWatch stopWatch = new StopWatch("rsa-verify");
        stopWatch.start();
        Signature signature = Signature.getInstance("SHA256WithRSA");
        signature.initVerify(getPublicKey(vName, keyStorePath, keyStorePassword));
        signature.update(data.getBytes(StandardCharsets.UTF_8));
        boolean verified = signature.verify(Base64.getDecoder()
                .decode(sign));
        stopWatch.stop();
        log.debug("rsa-verify: {}ms", stopWatch.getTotalTimeMillis());
        return verified;
    }

}
